import torch
from PIL import Image
from torchvision import transforms
from torch import nn
import torchxrayvision as xrv
from torch.utils.data import DataLoader

from conduct import test
from dataload import CovidCTDataset

model = xrv.models.DenseNet(num_classes=2, in_channels=3).cuda()
model.load_state_dict(torch.load(r'backup/DenseNet_medical_epoch400.pt'), False)
print('sucess')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
test_transformer = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    normalize
])

batchsize = 5
testset = CovidCTDataset(root_dir='data',
                             txt_COVID='data/testCT_COVID.txt',
                             txt_NonCOVID='data/testCT_NonCOVID.txt',
                             transform=test_transformer)
print(testset.__len__())
test_loader = DataLoader(testset, batch_size=batchsize, drop_last=False, shuffle=False)
torch.cuda.empty_cache()
targetlist, scorelist, predlist = test(model, test_loader)
print("predict", predlist)
print("targe", targetlist)